"""Package for reading and writing images, primarily using itk
This currently relies on my dicom package, but I want to replace this with
pydicom"""
import itk
import dicom
import os
from . import query
from . import ifilter
import numpy
import glob
import warnings

def orderDICOMFiles_old(dir=None,suffix="", verbose = False):
    try:
        if( dir == None ):
            return []
        filelist = glob.glob(dir+"/*"+suffix)
        dicom_files = []
        pos = {}
        for file in filelist:
            dfile = dicom.read_file(file, stop_before_pixels= True)
            if( verbose ): print "orderDICOMFiles: %s"%file
            if( not os.path.isdir(file ) ):
                try:
                    if( verbose ): print "orderDICOMFiles: InstanceNumber %s"%dfile.data_element('InstanceNumber').value
                    pos[dfile.data_element('InstanceNumber').value] = file
                except IOError:
                    print "not a dicom file",file
        orderlist = pos.keys()
        orderlist.sort()
        for o in orderlist:
            dicom_files.append(pos[o])
        if(verbose):
            print "ordered list of dicom files"
            for d in dicom_files:
                print d
        return dicom_files
    except Exception, error:
        print "Failed in orderDICOMFiles", error
        return []        
def orderDICOMFiles(dir=None,suffix=""):
	if( dir == None ):
	    return []
	filelist = glob.glob(dir+"/*"+suffix)
	dicom_files = {}
	for f in filelist:
	    try:
		df = dicom.read_file(f, stop_before_pixels = True)
		aqn = df.AcquisitionNumber
		inmbr = df.InstanceNumber
		files = dicom_files.get(aqn,[])
		files.append((inmbr,f))
		dicom_files[aqn] = files
	    except IOError:
		pass    
	    except InvalidDicomError:
		pass
	for aqn in dicom_files.keys():
	    files = dicom_files[aqn]
	    files.sort()
	    dicom_files[aqn] = files
	return dicom_files
def get_itk_image_from_dicom_series(dir=None, mode = "float", returnAcq=1 ):
    if( dir == None ):
        return None
    if( mode.lower() == "uchar"):
        img = itk.ImageSeriesReader.IUC3.New()
    elif( mode.lower() == 'sshort' ):
        img = itk.ImageSeriesReader.ISS3.New()
    elif( mode.lower() == 'ushort' ):
        img = itk.ImageSeriesReader.IUS3.New()
    else:
        img = itk.ImageSeriesReader.IF3.New()
    dfiles = orderDICOMFiles(dir=dir)[returnAcq]
    for dfile in dfiles:
        img.AddFileName(dfile[1])
    return img
        
def imageReader(fname,
              imgMode="float",
              dim = 3,
              iType = None, 
              returnITK = True, ):
    """use the ITK routines to read an image:
    
    fname: a file name or a directory containing DICOM images
    imgMode: string describing the pixel type to return "float", etc.
    dim: the dimension of the image (2 or 3)
    returnITK: return an itk image if True otherwise returns a numarray array"""
    try:
        if( os.path.isdir(fname) ): 
            print "reading dicom series for fixed image"
            fixedImageReader = get_itk_image_from_dicom_series(dir=fname,mode=imgMode)
        else:
            print "reading image from file",fname
            if (iType == None ):
                iitype = query.getITKType(imgMode)
                iType = itk.Image[iitype,dim]
            print "specified image type is",iType
            fixedImageReader = itk.ImageFileReader[iType].New(FileName=fname)
        fixedImageReader.Update()
        fixedImageReader.thisown = 0
        return fixedImageReader
    except Exception, error:
        print "failed in readImage", error

def readImage(fname,
              imgMode="float",
              dim = 3,
              iType = None, 
              returnITK = True, 
              verbose=False, 
              returnDescriptors = False):
    """use the ITK routines to read an image:
    
    fname: a file name or a directory containing DICOM images
    imgMode: string describing the pixel type to return "float", etc.
    dim: the dimension of the image (2 or 3)
    returnITK: return an itk image if True otherwise returns a numarray array"""
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        if (iType == None ):
            iitype = query.getITKType(imgMode)
            iType = itk.Image[iitype,dim]
        if( os.path.isdir(fname) ): 
            if(verbose):
                print "reading dicom series for fixed image"
            fixedImageReader = \
            get_itk_image_from_dicom_series(dir=fname,mode=imgMode)
        else:
            if(verbose):
                print "reading image from file",fname
            
                print "specified image type is",iType
            fixedImageReader = itk.ImageFileReader[iType].New(FileName=fname)
        img = fixedImageReader.GetOutput()
        if( returnITK ):
            img.Update()
            return img
        else:
            array = numpy.array(itk.PyBuffer[iType].GetArrayFromImage(fixedImageReader.GetOutput()))
            if( returnDescriptors):
                scl = query.get_itkVector(img.GetSpacing())
                origin = query.get_itkPoint(img.GetOrigin())
                direction = query.getITKDirection(img)
                return array, {"scale":scl, "origin":origin, "orientation":direction}
            return array

def numpyTypeToITKType(data):
    """compute the ITK data type corresponding to the numpy arrray data"""
    ntype = data.dtype.name
    dim = data.ndim
    if( ntype == 'uint8' ):
        return itk.Image[itk.UC,dim]
    elif( ntype == 'int8'):
        return itk.Image[itk.SC,dim]
    elif( ntype == 'uint16' ):
        return itk.Image[itk.US,dim]
    elif( ntype == 'int16' ):
        return itk.Image[itk.SS,dim]
    else:
        return itk.Image[itk.F,dim]
def writeImage(inputImg, fname="dummy.mha",
                outMode = None,
                debug=False,
                refimg = None ):
    """writes an ITK image. inMode, intype, and dim as argumetns now have no meaning if the
    argument passed is an ITK image. Instead the incoming image type and dimensions are grabbed
    from the images themselves. If outMode is none the same filetype as is input is assumed"""
    try:
        try:
            inType = query.numpyTypeToITKType(inputImg)
            img = itk.PyBuffer[intype].GetImageFromArray(inputImg)
        except:
            img = inputImg
            inType = query.getImageType(img)
        dim = inType.GetImageDimension()

        if( not outMode ):
            outType = inType
        else:
            otype = query.getITKType(outMode)
            outType = itk.Image[otype,dim]
        if( refimg ):
            print "setting header information to match input images"
            img = query.setHeaderData(img,refimg)
        if(debug): print "writeImage to %s"%fname
        writer = itk.ImageFileWriter[outType].New()
        if( inType != outType ):
            if( debug ): print "casting image to new output type"
            outputCast = ifilter.getCaster(img, outType=outType)
            oimg = outputCast.GetOutput()
            oimg.Update()          
            #printImageData(oimg,"Output image")  
            writer.SetInput( oimg )
        else:
            writer.SetInput(img)
        
        writer.SetFileName( fname )

        writer.Update()
    except Exception, error:
        print "failed in io.writeImage", error         
